import os
import numpy as np
import torch
import random
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from tools.my_dataset import DogCatDataset
from PIL import Image
from matplotlib import pyplot as plt


def set_seed(seed=1):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)


set_seed(1)  # 设置随机种子

# 参数设置
MAX_EPOCH = 1
BATCH_SIZE = 1
LR = 0.01
log_interval = 10
val_interval = 1
rmb_label = {"dog": 0, "cat": 1}


def transform_invert(img_, transform_train):
    """
    将data 进行反transfrom操作
    :param img_: tensor
    :param transform_train: torchvision.transforms
    :return: PIL image
    """
    if 'Normalize' in str(transform_train):
        norm_transform = list(filter(lambda x: isinstance(x, transforms.Normalize), transform_train.transforms))
        mean = torch.tensor(norm_transform[0].mean, dtype=img_.dtype, device=img_.device)
        std = torch.tensor(norm_transform[0].std, dtype=img_.dtype, device=img_.device)
        img_.mul_(std[:, None, None]).add_(mean[:, None, None])

    img_ = img_.transpose(0, 2).transpose(0, 1)  # C*H*W --> H*W*C
    img_ = np.array(img_) * 255

    if img_.shape[2] == 3:
        img_ = Image.fromarray(img_.astype('uint8')).convert('RGB')
    elif img_.shape[2] == 1:
        img_ = Image.fromarray(img_.astype('uint8').squeeze())
    else:
        raise Exception("Invalid img shape, expected 1 or 3 in axis 2, but got {}!".format(img_.shape[2]))

    return img_


# ============================ step 1/5 数据 ============================
split_dir = os.path.join("..", "data", "dog_cat_split")
train_dir = os.path.join(split_dir, "train")
valid_dir = os.path.join(split_dir, "valid")

norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]
t = torch.tensor([224, 224])
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),

    # 1 CenterCrop
    # transforms.CenterCrop(100),  # 512

    # 2 RandomCrop
    # transforms.RandomCrop(150, padding=16),
    # transforms.RandomCrop(224, padding=(16, 64)),
    # transforms.RandomCrop(224, padding=16, fill=(255, 0, 0)),
    # transforms.RandomCrop(512, pad_if_needed=True),   # 填充模式下，当裁剪大小大于原图时，pad_if_needed=True
    # transforms.RandomCrop(224, padding=64, padding_mode='edge'),
    # transforms.RandomCrop(224, padding=64, padding_mode='reflect'),   # 镜像模式
    # transforms.RandomCrop(1024, padding=1024, padding_mode='symmetric'),

    # 3 RandomResizedCrop 随机大小、长宽比裁剪
    # transforms.RandomResizedCrop(size=224, scale=(0.5, 0.5), interpolation=Image.NEAREST),

    # 4 FiveCrop
    # 左下角、左上角、右下角、右上角、中心，进行裁剪，得到5张图片
    # 返回一个tuple
    # transforms.FiveCrop(112),
    # 用户自定义方法
    # transforms.Lambda(lambda crops: torch.stack([(transforms.ToTensor()(crop)) for crop in crops])),

    # 5 TenCrop
    # 左下角、左上角、右下角、右上角、中心，进行裁剪，并进行翻转，得到10张图片
    # transforms.TenCrop(112, vertical_flip=False),
    # transforms.Lambda(lambda crops: torch.stack([(transforms.ToTensor()(crop)) for crop in crops])),

    # 1 Horizontal Flip
    # 依概率p水平翻转
    # transforms.RandomHorizontalFlip(p=1),

    # 2 Vertical Flip
    # 依概率p垂直翻转
    # transforms.RandomVerticalFlip(p=0.5),

    # 3 RandomRotation
    # transforms.RandomRotation(90),
    # transforms.RandomRotation((90), expand=True),
    # transforms.RandomRotation(30, center=(0, 0)),
    # transforms.RandomRotation(30, center=(112, 112), expand=True),   # expand only for center rotation

    # 1 Pad
    # 填充
    # transforms.Pad(padding=32, fill=(200, 0, 0), padding_mode='constant'),    # 填充模式
    # transforms.Pad(padding=(8, 64), fill=(255, 0, 0), padding_mode='constant'),
    # transforms.Pad(padding=(8, 16, 32, 64), fill=(255, 0, 0), padding_mode='constant'),
    # transforms.Pad(padding=(8, 16, 32, 64), fill=(255, 0, 0), padding_mode='symmetric'),  # 镜像模式

    # ColorJitter(brightness=0,   # 亮度
    #             contrast=0,     # 对比度
    #             saturation=0,   # 饱和度
    #             hue=0)          # 色相
    # 调整高亮、对比度、饱和度和色相
    # brightness、contrast、saturation    为一个数a时，在[max(0,1-a), 1+a]中随机选择，当为(a,b)时，在[a,b]中随机选择
    # hue 为一个数a时，在[-a,a]中随机选择,0<=a<=0.5，当为(a,b)时，在[a,b]中随机选择,-0.5<=a<=b<=0.5
    # transforms.ColorJitter(brightness=0.5),
    # transforms.ColorJitter(contrast=0.1),
    # transforms.ColorJitter(saturation=0.1),
    # transforms.ColorJitter(hue=0.5),

    # 3 Grayscale
    # 转灰度图
    # transforms.Grayscale(num_output_channels=3),

    # RandomGrayscale
    # 依概率p转灰度图
    # num_output_channels 只能设为1或3
    # p 是图片转为灰度图的概率
    # transforms.RandomGrayscale(num_output_channels=3, p=0.1),

    # transforms.LinearTransformation(t,t),

    # 仿射变换，是二维线性变换，由五种基本原子变换构成，分别是旋转、平移、缩放、错切、旋转
    # RandomAffine(degrees,   # 旋转角度
    #              translate=None,    # 平移区间(a,b), -img_width * a < dx < img_width * a
    #              scale=None,    # 缩放比例，以面积为单位（0,1）之间
    #              shear=None,    # 错切，为a时，仅在x错切，角度在(-a,a)之间；为(a,b)，a设置x，b设置y；为(a,b,c,d)，a,b设置x，c,d设置y
    #              resample=False,    # 重采样方式，NEAREST,NILINEAR,BICUBIC
    #              fillcolor=0)   # 填充颜色，默认黑色
    # transforms.RandomAffine(degrees=0),
    # transforms.RandomAffine(degrees=0, translate=(0.2, 0.2), fillcolor=(0, 0, 0)),
    # transforms.RandomAffine(degrees=0, scale=(0.7, 0.7)),
    # transforms.RandomAffine(degrees=0, shear=(0, 0, 0, 45)),
    # transforms.RandomAffine(degrees=0, shear=90, fillcolor=(255, 0, 0)),

    # 随机遮挡
    # RandomErasing(p=0.5,    # 概率值
    #               scale=(0.02,0.33),    # 遮挡区域面积
    #               ratio=(0.3,3.3),      # 遮挡区域长宽比
    #               value=0,      # 遮挡区域像素值 (R,G,B) or (gray) 或任意字符串
    #               inplace=False)    # 原位操作
    # transforms.ToTensor(),
    # transforms.RandomErasing(p=1, scale=(0.02, 0.33), ratio=(0.3, 3.3), value=(255, 0, 0)),
    # transforms.RandomErasing(p=0.5, scale= (0.01, 0.33), ratio=(0.4, 2.3), value='liijl'),

    # 1 RandomChoice
    # 从给定的一系列transforms中选一个进行操作
    # transforms.RandomChoice([transforms.RandomVerticalFlip(p=1), transforms.RandomHorizontalFlip(p=1)]),

    # 2 RandomApply
    # 给一个transform加上概率，依概率进行选择操作
    # transforms.RandomApply([transforms.RandomAffine(degrees=0, shear=45, fillcolor=(255, 0, 0)),
    #                         transforms.Grayscale(num_output_channels=3)], p=0.5),
    # 3 RandomOrder
    # 将transforms中的操作随机打乱
    # transforms.RandomOrder([transforms.RandomRotation(15),
    #                         transforms.Pad(padding=32),
    #                         transforms.RandomAffine(degrees=0, translate=(0.01, 0.1), scale=(0.9, 1.1))]),

    transforms.ToTensor(),
    # transforms.Normalize(norm_mean, norm_std),
])

valid_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std)
])

# 构建MyDataset实例
train_data = DogCatDataset(data_dir=train_dir, transform=train_transform)
valid_data = DogCatDataset(data_dir=valid_dir, transform=valid_transform)

# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)

# ============================ step 5/5 训练 ============================
for epoch in range(MAX_EPOCH):
    for i, data in enumerate(train_loader):
        inputs, labels = data  # B C H W
        img_tensor = inputs[0, ...]  # C H W
        # 将张量逆变换为PIL img，用于可视化
        img = transform_invert(img_tensor, train_transform)
        plt.imshow(img)
        # plt.savefig("/home/vistar/Desktop/img/17.png", bbox_inches='tight')
        plt.show()
        plt.pause(0.5)
        plt.close()
        break

        # bs, ncrops, c, h, w = inputs.shape
        # for n in range(ncrops):
        #     img_tensor = inputs[0, n, ...]  # C H W
        #     img = transform_invert(img_tensor, train_transform)
        #     plt.imshow(img)
        #     # plt.savefig("/home/vistar/Desktop/"+str(n)+".png", bbox_inches='tight')
        #     plt.show()
        #     plt.pause(1)
        # break
